# -*- coding: utf-8 -*-

"""
Unicorn的实时地图投射工具
========================

功能:
1.  截取并选择屏幕上的任意两个区域：游戏地图区域和投射区域。
2.  实时捕捉地图区域的画面，可选MSS或DXGI方式。
3.  通过HSV颜色空间调整，过滤掉不需要的背景，只保留地图核心内容。
4.  将处理后的画面实时、置顶、透明、可穿透点击地投射到指定区域。
5.  所有配置均可随时调整。
6.  可以保存和读取配置，方便快速切换。

使用注意:
- 请确保 unicorn.jpg 图片文件与此脚本位于同一目录下。它将作为程序图标和背景。

依赖库:
- PyQt6: 用于构建图形用户界面 (GUI)。
- mss: 用于高效的实时屏幕捕获。
- d3dshot: 用于高性能的DXGI屏幕捕射。
- OpenCV-Python: 用于图像处理（颜色空间转换和蒙版）。
- NumPy: OpenCV 的依赖，用于处理图像数据。
- pywin32: 用于实现窗口的“点击穿透”效果 (可选，PyQt6自带功能通常足够)。

安装命令:
# 请按顺序执行以下三条命令
pip install --no-deps d3dshot
pip install --upgrade Pillow
pip install --upgrade comtypes
# 再安装其他库
pip install PyQt6 mss opencv-python numpy pywin32
"""

import sys
import os
import mss
import numpy as np
import cv2
import json
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QPushButton, QLabel, QSlider,
    QVBoxLayout, QHBoxLayout, QGridLayout, QGroupBox, QSpinBox, QFileDialog, QMessageBox,
    QRadioButton
)
from PyQt6.QtCore import Qt, QTimer, QRect, QPoint, QSize
from PyQt6.QtGui import (
    QPainter, QPen, QBrush, QPixmap, QImage, QScreen, QFont, QColor, QPainterPath, QIcon
)

# 尝试导入d3dshot
try:
    import d3dshot
    D3DSHOT_AVAILABLE = True
except ImportError:
    D3DSHOT_AVAILABLE = False
    print("警告: d3dshot 库未安装，DXGI捕获方式不可用。请按说明分步安装。")

# 尝试导入pywin32
try:
    import win32gui
    import win32con
    PYWIN32_AVAILABLE = True
except ImportError:
    PYWIN32_AVAILABLE = False
    print("警告: pywin32 库未安装，点击穿透功能可能受限。建议安装: pip install pywin32")


def resource_path(relative_path):
    """ 获取资源的绝对路径, 适用于开发环境和PyInstaller打包环境 """
    try:
        # PyInstaller 创建一个临时文件夹, 并把路径存储在 _MEIPASS 中
        base_path = sys._MEIPASS
    except Exception:
        base_path = os.path.abspath(".")

    return os.path.join(base_path, relative_path)


class SelectionWindow(QWidget):
    """
    一个全屏的、半透明的窗口，用于让用户框选区域。
    """
    def __init__(self, parent=None):
        super().__init__(parent)
        self.parent = parent
        self.setWindowFlags(
            Qt.WindowType.FramelessWindowHint |
            Qt.WindowType.WindowStaysOnTopHint |
            Qt.WindowType.Tool
        )
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)
        
        screen = QApplication.primaryScreen()
        self.background_pixmap = screen.grabWindow(0)
        self.setGeometry(screen.geometry())

        self.start_point = QPoint()
        self.end_point = QPoint()
        self.current_rect = QRect()
        self.selecting = False
        self.state = "map"
        self.map_rect = None
        self.projection_rect = None
        
        self.info_label = QLabel(self)
        self.info_label.setFont(QFont("Microsoft YaHei", 24, QFont.Weight.Bold))
        self.info_label.setStyleSheet("color: white; background-color: rgba(0, 0, 0, 150); padding: 10px; border-radius: 5px;")
        self.info_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.update_info_text()

    def update_info_text(self):
        """ 更新提示用户的文本, 并加入退出提示 """
        base_text = ""
        if self.state == "map":
            base_text = "请用鼠标拖拽，框选出【游戏地图】区域"
        else:
            base_text = "请用鼠标拖拽，框选出【投射显示】区域"
        
        self.info_label.setText(f"{base_text}\n(按 Esc 键退出)")
        self.info_label.adjustSize()
        self.info_label.move(int((self.width() - self.info_label.width()) / 2), 50)

    def keyPressEvent(self, event):
        """ 监听键盘事件, 实现Esc键退出 """
        if event.key() == Qt.Key.Key_Escape:
            self.parent.show() # 显示主窗口
            self.close()

    def mousePressEvent(self, event):
        if event.button() == Qt.MouseButton.LeftButton:
            self.selecting = True
            self.start_point = event.pos()
            self.end_point = event.pos()
            self.update()

    def mouseMoveEvent(self, event):
        if self.selecting:
            self.end_point = event.pos()
            self.current_rect = QRect(self.start_point, self.end_point).normalized()
            self.update()

    def mouseReleaseEvent(self, event):
        if event.button() == Qt.MouseButton.LeftButton and self.selecting:
            self.selecting = False
            final_rect = QRect(self.start_point, self.end_point).normalized()
            
            if self.state == "map":
                self.map_rect = final_rect
                self.state = "projection"
                self.update_info_text()
            else:
                self.projection_rect = final_rect
                self.parent.on_selection_complete(self.map_rect, self.projection_rect)
                self.close()
            self.update()

    def paintEvent(self, event):
        painter = QPainter(self)
        painter.setRenderHint(QPainter.RenderHint.Antialiasing)
        painter.drawPixmap(self.rect(), self.background_pixmap)
        
        path = QPainterPath()
        path.addRect(self.rect().toRectF())
        if self.map_rect:
            path.addRect(self.map_rect.toRectF())
        if self.selecting:
            path.addRect(self.current_rect.toRectF())
        
        painter.fillPath(path, QBrush(QColor(0, 0, 0, 128)))

        if self.map_rect:
            pen = QPen(QColor(255, 0, 0, 200), 3, Qt.PenStyle.SolidLine)
            painter.setPen(pen)
            painter.setBrush(Qt.BrushStyle.NoBrush)
            painter.drawRect(self.map_rect)
            painter.drawText(self.map_rect.topLeft() - QPoint(0, 5), "地图区域")

        if self.selecting:
            pen_color = QColor(255, 0, 0, 255) if self.state == "map" else QColor(0, 128, 255, 255)
            pen = QPen(pen_color, 2, Qt.PenStyle.DashLine)
            painter.setPen(pen)
            brush_color = pen_color.lighter(180)
            brush_color.setAlphaF(0.3)
            painter.setBrush(QBrush(brush_color))
            painter.drawRect(self.current_rect)

class ProjectionWindow(QWidget):
    def __init__(self):
        super().__init__()
        self.pixmap_to_draw = None
        self.setWindowFlags(
            Qt.WindowType.FramelessWindowHint |
            Qt.WindowType.WindowStaysOnTopHint |
            Qt.WindowType.Tool
        )
        self.setAttribute(Qt.WidgetAttribute.WA_TranslucentBackground)
        self.setAttribute(Qt.WidgetAttribute.WA_TransparentForMouseEvents)

    def set_pixmap(self, pixmap):
        self.pixmap_to_draw = pixmap
        self.update()

    def paintEvent(self, event):
        if self.pixmap_to_draw:
            painter = QPainter(self)
            painter.drawPixmap(self.rect(), self.pixmap_to_draw)
            
    def showEvent(self, event):
        super().showEvent(event)
        if PYWIN32_AVAILABLE:
            try:
                hwnd = self.winId()
                ex_style = win32gui.GetWindowLong(hwnd, win32con.GWL_EXSTYLE)
                win32gui.SetWindowLong(hwnd, win32con.GWL_EXSTYLE, ex_style | win32con.WS_EX_TRANSPARENT)
            except Exception as e:
                print(f"设置点击穿透时出错: {e}")

class MainController(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("Unicorn的实时地图投射工具 - 控制面板")
        self.setFont(QFont("Microsoft YaHei", 10))
        self.setFixedSize(680, 650)
        
        self.icon_path = resource_path("unicorn.jpg")
        self.setWindowIcon(QIcon(self.icon_path))

        self.map_rect = None
        self.projection_rect = None
        self.map_screenshot = None
        self.is_projecting = False

        self.selection_win = None
        self.projection_win = ProjectionWindow()
        self.capture_timer = QTimer(self)
        self.capture_timer.setInterval(16)
        self.capture_timer.timeout.connect(self.update_projection_frame)
        
        self.sct = mss.mss()
        self.d_shot = None
        if D3DSHOT_AVAILABLE:
            try:
                self.d_shot = d3dshot.create(capture_output="numpy")
            except Exception as e:
                print(f"初始化DXGI失败: {e}")
                self.d_shot = None
        self.capture_method = "mss"

        self.init_ui()

    def init_ui(self):
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        
        stylesheet = f"""
            QMainWindow {{
                border-image: url({self.icon_path.replace(os.sep, '/')});
            }}
            QGroupBox {{
                background-color: rgba(255, 255, 255, 190);
                border: 1px solid #CCCCCC; border-radius: 8px;
                margin-top: 1em; font-weight: bold;
            }}
            QGroupBox::title {{
                subcontrol-origin: margin; subcontrol-position: top center; padding: 0 10px;
            }}
            QLabel, QRadioButton {{ background-color: transparent; font-weight: bold; }}
        """
        self.setStyleSheet(stylesheet)

        main_layout = QVBoxLayout(central_widget)
        main_layout.setContentsMargins(20, 20, 20, 20)

        # 1. 配置区域
        config_group = QGroupBox("核心配置")
        config_layout = QVBoxLayout()
        self.select_button = QPushButton("1. 选取地图与投射区域")
        self.select_button.setStyleSheet("padding: 10px; background-color: #007BFF; color: white; border-radius: 5px; font-weight: bold;")
        self.select_button.clicked.connect(self.start_selection)
        config_layout.addWidget(self.select_button)
        
        size_layout = QHBoxLayout()
        self.size_spinbox = QSpinBox()
        self.size_spinbox.setRange(50, 800)
        self.size_spinbox.setValue(200)
        self.size_spinbox.setSuffix(" 像素 (宽)")
        self.size_spinbox.valueChanged.connect(self.update_projection_size)
        self.size_preview_label = QLabel("示例: 200x200")
        size_layout.addWidget(QLabel("投射尺寸:"))
        size_layout.addWidget(self.size_spinbox)
        size_layout.addWidget(self.size_preview_label)
        config_layout.addLayout(size_layout)
        config_group.setLayout(config_layout)
        main_layout.addWidget(config_group)

        # 2. 捕获方式区域
        capture_group = QGroupBox("捕获方式")
        capture_layout = QHBoxLayout()
        self.mss_radio = QRadioButton("MSS (兼容性好)")
        self.dxgi_radio = QRadioButton("DXGI (性能更高)")
        self.mss_radio.setChecked(True)
        if not D3DSHOT_AVAILABLE or self.d_shot is None:
            self.dxgi_radio.setEnabled(False)
            self.dxgi_radio.setText("DXGI (不可用)")
        self.mss_radio.toggled.connect(lambda: self.set_capture_method("mss"))
        self.dxgi_radio.toggled.connect(lambda: self.set_capture_method("dxgi"))
        capture_layout.addWidget(self.mss_radio)
        capture_layout.addWidget(self.dxgi_radio)
        capture_group.setLayout(capture_layout)
        main_layout.addWidget(capture_group)

        # 3. 配置管理区域
        config_manage_group = QGroupBox("配置管理")
        config_manage_layout = QHBoxLayout()
        save_button = QPushButton("保存配置")
        load_button = QPushButton("读取配置")
        save_button.setStyleSheet("padding: 8px; background-color: #17A2B8; color: white; border-radius: 5px; font-weight: bold;")
        load_button.setStyleSheet("padding: 8px; background-color: #6C757D; color: white; border-radius: 5px; font-weight: bold;")
        save_button.clicked.connect(self.save_config)
        load_button.clicked.connect(self.load_config)
        config_manage_layout.addWidget(save_button)
        config_manage_layout.addWidget(load_button)
        config_manage_group.setLayout(config_manage_layout)
        main_layout.addWidget(config_manage_group)

        # 4. 图像处理区域
        processing_group = QGroupBox("实时图像处理 (HSV颜色过滤)")
        processing_layout = QGridLayout()
        self.preview_label = QLabel("请先选择地图区域以预览效果")
        self.preview_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.preview_label.setFixedSize(300, 200)
        self.preview_label.setStyleSheet("background-color: #333; color: white; border: 1px solid #555; border-radius: 5px;")
        processing_layout.addWidget(self.preview_label, 0, 0, 4, 1)

        self.sliders = {}
        slider_definitions = [
            ("h_max", "H Max", 0, 179, 179), ("h_min", "H Min", 0, 179, 0),
            ("s_max", "S Max", 0, 255, 255), ("s_min", "S Min", 0, 255, 0),
            ("v_max", "V Max", 0, 255, 255), ("v_min", "V Min", 0, 255, 0),
        ]

        for i, (key, name, min_v, max_v, def_v) in enumerate(slider_definitions):
            slider, label = self._create_slider(name, min_v, max_v, def_v)
            self.sliders[key] = slider
            slider.valueChanged.connect(self.update_hsv_preview)
            col = (i // 2) + 1
            row_for_label = (i % 2) * 2
            row_for_slider = row_for_label + 1
            processing_layout.addWidget(label, row_for_label, col)
            processing_layout.addWidget(slider, row_for_slider, col)
        
        processing_layout.setColumnStretch(0, 3)
        processing_layout.setColumnStretch(1, 1)
        processing_layout.setColumnStretch(2, 1)
        processing_layout.setColumnStretch(3, 1)
        processing_group.setLayout(processing_layout)
        main_layout.addWidget(processing_group)

        # 5. 运行控制区域
        run_group = QGroupBox("运行控制")
        run_layout = QHBoxLayout()
        self.run_button = QPushButton("开始投射")
        self.run_button.setCheckable(True)
        self.run_button.setStyleSheet("""
            QPushButton { padding: 10px; background-color: #28A745; color: white; border-radius: 5px; font-weight: bold; }
            QPushButton:checked { background-color: #DC3545; }
        """)
        self.run_button.clicked.connect(self.toggle_projection)
        run_layout.addWidget(self.run_button)
        run_group.setLayout(run_layout)
        main_layout.addWidget(run_group)

    def set_capture_method(self, method):
        if method == "mss":
            self.capture_method = "mss"
        elif method == "dxgi" and D3DSHOT_AVAILABLE:
            self.capture_method = "dxgi"

    def _create_slider(self, name, min_val, max_val, default_val):
        label = QLabel(f"{name}: {default_val}")
        slider = QSlider(Qt.Orientation.Vertical)
        slider.setRange(min_val, max_val)
        slider.setValue(default_val)
        slider.setMinimumHeight(80)
        slider.valueChanged.connect(lambda val, l=label, n=name: l.setText(f"{n}: {val}"))
        return slider, label

    def start_selection(self):
        self.hide() 
        QTimer.singleShot(100, self._show_selection_window)

    def _show_selection_window(self):
        self.selection_win = SelectionWindow(self)
        self.selection_win.show()
        self.selection_win.activateWindow()
        self.selection_win.raise_()

    def on_selection_complete(self, map_rect, projection_rect):
        self.map_rect = map_rect
        self.projection_rect = projection_rect
        self._update_map_screenshot()
        self.update_hsv_preview()
        self.update_projection_size()
        self.show() 
        self.activateWindow()
        self.raise_()

    def _update_map_screenshot(self):
        if not self.map_rect: return
        screen = QApplication.primaryScreen()
        pixmap = screen.grabWindow(0, self.map_rect.x(), self.map_rect.y(), self.map_rect.width(), self.map_rect.height())
        self.map_screenshot = self.qpixmap_to_cv_image(pixmap)

    def update_hsv_preview(self):
        if self.map_screenshot is None: return
        processed_img = self.process_image(self.map_screenshot)
        q_img = self.cv_image_to_qimage(processed_img)
        pixmap = QPixmap.fromImage(q_img)
        self.preview_label.setPixmap(pixmap.scaled(
            self.preview_label.size(), 
            Qt.AspectRatioMode.KeepAspectRatio, 
            Qt.TransformationMode.SmoothTransformation
        ))

    def update_projection_size(self):
        if self.projection_rect is None or self.map_rect is None or self.map_rect.width() == 0: return
        width = self.size_spinbox.value()
        aspect_ratio = self.map_rect.height() / self.map_rect.width()
        height = int(width * aspect_ratio)
        self.projection_win.setGeometry(self.projection_rect.x(), self.projection_rect.y(), width, height)
        self.size_preview_label.setText(f"示例: {width}x{height}")

    def toggle_projection(self, checked):
        if checked:
            if self.map_rect is None or self.projection_rect is None:
                QMessageBox.warning(self, "错误", "请先选择地图和投射区域！")
                self.run_button.setChecked(False)
                return
            
            if self.capture_method == "dxgi" and self.d_shot is None:
                QMessageBox.critical(self, "DXGI错误", "DXGI 捕获器初始化失败，无法开始。")
                self.run_button.setChecked(False)
                return

            self.is_projecting = True
            self.run_button.setText("停止投射")
            self.projection_win.show()
            self.capture_timer.start()
        else:
            self.is_projecting = False
            self.run_button.setText("开始投射")
            self.capture_timer.stop()
            self.projection_win.hide()
            
    def update_projection_frame(self):
        if not self.is_projecting: return
        
        img = None
        x, y, w, h = self.map_rect.getRect()

        if self.capture_method == "mss":
            mon = {"top": y, "left": x, "width": w, "height": h}
            img_bgra = self.sct.grab(mon)
            img = np.array(img_bgra)
        elif self.capture_method == "dxgi" and self.d_shot:
            try:
                region = (x, y, x + w, y + h)
                img_rgb = self.d_shot.screenshot(region=region)
                if img_rgb is not None:
                    img = cv2.cvtColor(img_rgb, cv2.COLOR_RGB2BGRA)
            except Exception as e:
                print(f"DXGI 捕获帧失败: {e}")
                return

        if img is None or img.size == 0: return

        processed_img = self.process_image(img)
        q_img = self.cv_image_to_qimage(processed_img)
        pixmap = QPixmap.fromImage(q_img)
        self.projection_win.set_pixmap(pixmap)
        
    def process_image(self, img_bgra):
        img_bgr = cv2.cvtColor(img_bgra, cv2.COLOR_BGRA2BGR)
        hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
        
        hsv_values = {key: slider.value() for key, slider in self.sliders.items()}
        lower_bound = np.array([hsv_values["h_min"], hsv_values["s_min"], hsv_values["v_min"]])
        upper_bound = np.array([hsv_values["h_max"], hsv_values["s_max"], hsv_values["v_max"]])
        mask = cv2.inRange(hsv, lower_bound, upper_bound)

        result_bgra = np.zeros_like(img_bgra)
        result_bgra[mask > 0] = img_bgra[mask > 0]
        return result_bgra

    def save_config(self):
        if not self.map_rect:
            QMessageBox.warning(self, "提示", "没有可保存的配置。请先选择区域。")
            return

        filePath, _ = QFileDialog.getSaveFileName(self, "保存配置", "", "JSON Files (*.json)")
        if not filePath: return

        config = {
            "map_rect": self.map_rect.getRect(),
            "projection_rect": self.projection_rect.getRect(),
            "projection_size": self.size_spinbox.value(),
            "hsv_values": {key: slider.value() for key, slider in self.sliders.items()},
            "capture_method": self.capture_method
        }
        try:
            with open(filePath, 'w') as f:
                json.dump(config, f, indent=4)
            QMessageBox.information(self, "成功", f"配置已保存到:\n{filePath}")
        except Exception as e:
            QMessageBox.critical(self, "错误", f"保存配置失败: {e}")

    def load_config(self):
        filePath, _ = QFileDialog.getOpenFileName(self, "读取配置", "", "JSON Files (*.json)")
        if not filePath: return

        try:
            with open(filePath, 'r') as f:
                config = json.load(f)

            self.map_rect = QRect(*config["map_rect"])
            self.projection_rect = QRect(*config["projection_rect"])
            self.size_spinbox.setValue(config["projection_size"])

            for key, value in config["hsv_values"].items():
                if key in self.sliders:
                    self.sliders[key].setValue(value)
            
            capture_method = config.get("capture_method", "mss")
            if capture_method == "dxgi" and D3DSHOT_AVAILABLE and self.d_shot is not None:
                self.dxgi_radio.setChecked(True)
            else:
                self.mss_radio.setChecked(True)
            self.set_capture_method(capture_method)

            self._update_map_screenshot()
            self.update_hsv_preview()
            self.update_projection_size()
            QMessageBox.information(self, "成功", f"已成功加载配置:\n{filePath}")

        except Exception as e:
            QMessageBox.critical(self, "错误", f"加载配置失败: {e}")

    def qpixmap_to_cv_image(self, pixmap):
        qimage = pixmap.toImage().convertToFormat(QImage.Format.Format_ARGB32)
        ptr = qimage.bits()
        ptr.setsize(qimage.sizeInBytes())
        arr = np.frombuffer(ptr, np.uint8).reshape((qimage.height(), qimage.width(), 4))
        return arr.copy()

    def cv_image_to_qimage(self, cv_img):
        h, w, ch = cv_img.shape
        bytes_per_line = ch * w
        return QImage(cv_img.data, w, h, bytes_per_line, QImage.Format.Format_ARGB32)

    def closeEvent(self, event):
        self.capture_timer.stop()
        if self.d_shot:
             try:
                # d3dshot does not have a formal stop method, it's handled by garbage collection
                pass
             except Exception as e:
                print(f"停止DXGI时出错: {e}")
        if self.projection_win: self.projection_win.close()
        if self.selection_win: self.selection_win.close()
        event.accept()

if __name__ == '__main__':
    app = QApplication(sys.argv)
    controller = MainController()
    controller.show()
    sys.exit(app.exec())
